"""
远程请求数据传输后端

date: 2025/10/19
author: SiHeng Tang
file: rrtlib.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
import json
import logging
import os
import socket
import time
import zlib
from typing import Literal, Any

import serial

logger = logging.getLogger(__name__)


class _BaseRemoteRequestTelemetry(object):
    """
    基础发送后端，
    MJ-RPC 使用变长数据包，CRC32 校验，zlib 压缩和重传机制实现传输层，并不绝对完全可靠，作者并不精通于此

    典型的数据包传输流程::

        | rpc_object | => | encoded_json_str | => | zlib_compress_data | => | CRC32 | binary_data |
        = SPLIT =>
        | current_index | total_count | chunk_1 |
        ...
        | current_index | total_count | chunk_x |
        = TX/RX =>
        = RECV AND CHECK EVERY CRC32 AND GIVE RESPONSE =>
        | raw_data |  => | uncompress_json_data | => | rpc_json_data |

    二进制传输没有 json 的预处理和后处理，更加简化::

        | binary_data | => | CRC32 | binary_data |
        = SPLIT =>
        | current_index | total_count | chunk_1 |
        ...
        | current_index | total_count | chunk_x |
        = TX/RX =>
        = RECV AND CHECK EVERY CRC32 AND GIVE RESPONSE =>
        | binary_data |
    """
    ACK_RESPONSE = b'\x9d'
    FRAME_LENGTH = 4
    CRC32_LENGTH = 4
    PACK_INDEX_LENGTH = 2
    ACK_LENGTH = len(ACK_RESPONSE)

    def __init__(self,
                 max_chunk_length: int = 1024,
                 chunk_interval: float = 0.0,
                 zlib_compress_level: int = 6,
                 code_order: Literal["little", "big"] = "little",
                 json_str_coding: str = "utf-8",
                 pack_loss_resend_cnt: int = 5,
                 hand_shake_timeout: float = 65535.0,
                 transmit_timeout: float = 10.0,
                 reconnect_timeout: float = 15.0,
                 sending_mode: Literal["bin", "json"] = "bin"):
        """
        初始化连接，RRT 后端最核心的参数是延时发送和包长度切分，必须保证这些参数正确设置，
        默认的参数为使用网桥通信的飞行器优化
        :param max_chunk_length: 最大单个数据包长度，默认 1024
        :param chunk_interval: 数据包发送延时，网络协议往往不需要，默认 0
        :param zlib_compress_level: zlib 数据包压缩级别，默认 6
        :param code_order: 编码字节大小端，默认小段
        :param json_str_coding: json 字符串化对象编码方式，默认 utf-8
        :param pack_loss_resend_cnt: 数据包丢失后重传次数，默认5次
        :param hand_shake_timeout: 握手前等待时间，默认 65535 秒
        :param transmit_timeout: 传输时超时，默认1秒
        :param reconnect_timeout: 重新连接时超时，默认3秒
        :param sending_mode: 发送模式，可以选择二进制或者 RPC json 发送模式
        """
        self.chunk_interval = chunk_interval
        self.max_chunk_length = max_chunk_length

        self.code_order = code_order
        self.json_str_coding = json_str_coding
        self.zlib_compress_level = zlib_compress_level

        self.pack_loss_reconnect_cnt = pack_loss_resend_cnt
        self.transmit_timeout = transmit_timeout
        self.hand_shake_timeout = hand_shake_timeout
        self.reconnect_timeout = reconnect_timeout

        self.current_link_out_speed: float = 0.0
        self.current_link_in_speed: float = 0.0

        self.is_connected = False
        self.retry_times = 0
        self.sending_mode: Literal["bin", "json"] = sending_mode

    @property
    def speed(self):
        """输入数据流和输出数据流分别的速度"""
        return self.current_link_in_speed, self.current_link_out_speed

    def send_request_obj(self, req_obj: dict | list):
        """
        发送远程请求对象
        :param req_obj: 远程请求对象
        """
        start_time = time.time()
        transferred_bytes_count = 0

        data = json.dumps(req_obj, separators=(",", ":"), indent=None).encode(self.json_str_coding)
        data = zlib.compress(data, level=self.zlib_compress_level)

        data = zlib.crc32(data).to_bytes(self.CRC32_LENGTH, self.code_order) + data
        view = memoryview(data)

        pack_cnt = len(view) // self.max_chunk_length + int(bool(len(view) % self.max_chunk_length))
        pack_idx_current = 1

        # 发送和验证逻辑，从发送缓冲区读取不超过最大长度的字节，打包为符合格式的数据包，通过发送函数发送
        # 检查函数是否完全发送数值，是否按时返回确认，出错就递增重传次数，超过限制直接报错
        while view:
            if self.retry_times >= self.pack_loss_reconnect_cnt:
                raise ConnectionError("Data can not reach destination.")
            if self.retry_times > 0:
                timeout = self.reconnect_timeout
            else:
                timeout = self.transmit_timeout

            chunk = bytes(view[:self.max_chunk_length])
            payload = (pack_idx_current.to_bytes(self.PACK_INDEX_LENGTH, self.code_order) +
                       pack_cnt.to_bytes(self.PACK_INDEX_LENGTH, self.code_order) +
                       len(chunk).to_bytes(self.FRAME_LENGTH, self.code_order) +
                       chunk)

            sent_length = self._send_chunk(payload)
            if sent_length != len(payload):
                logger.error(f"sender can not send expected length {len(payload)} only {sent_length}, retrying")
                self.retry_times += 1
                continue

            try:
                ack_bit = self._recv_chunk(self.ACK_LENGTH, timeout)
            except (ConnectionError, TimeoutError) as err:
                logger.error(f"data transmit error: {str(err)}, retrying")
            else:
                if ack_bit != self.ACK_RESPONSE:
                    logger.error(f"data transmit error: ack response invalid, resending data.")
                    self.retry_times += 1
                    continue

            view = view[len(chunk):]

            self.retry_times = 0
            pack_idx_current += 1
            transferred_bytes_count += sent_length

            if self.chunk_interval > 0.1:
                time.sleep(self.chunk_interval)

        self.current_link_out_speed = transferred_bytes_count / (time.time() - start_time)

    def receive_response_obj(self) -> Any | None:
        """
        接收远程请求对象
        :raises ConnectionError: 超时或数据包校验失败，并且重试失败
        """
        start_time = time.time()

        if not self.is_connected:
            timeout = self.hand_shake_timeout
        elif self.retry_times > 0:
            timeout = self.reconnect_timeout
        else:
            timeout = self.transmit_timeout

        picked_raw_data = bytes()
        current_pack_idx = 1

        # 接收数据包，按照编码格式，需要根据数据包头判断当前包序号是否正确递增，同时执行 CRC32 校验，成功后回报，
        while True:
            if self.retry_times >= self.pack_loss_reconnect_cnt:
                raise ConnectionError("run out retry times, can not parse data correctly.")

            try:
                pack_header = self._recv_chunk(self.PACK_INDEX_LENGTH * 2 + self.FRAME_LENGTH, timeout)
                pack_index = int.from_bytes(pack_header[:self.PACK_INDEX_LENGTH], self.code_order)
                pack_cnt = int.from_bytes(pack_header[self.PACK_INDEX_LENGTH:self.PACK_INDEX_LENGTH * 2],
                                          self.code_order)
                chunk_length = int.from_bytes(pack_header[self.PACK_INDEX_LENGTH * 2:], self.code_order)

                if pack_index > pack_cnt or current_pack_idx != pack_index:
                    logger.error(f"error pack index {pack_index} in all pack count {pack_cnt}, retrying")
                    self.retry_times += 1
                    continue
            except TimeoutError as err:
                raise ConnectionError(f"no connect in ready time, {str(err)}")
            except (ConnectionError, OSError) as err:
                raise ConnectionError(f"receive data error when prase header, {str(err)}")

            try:
                raw_pack = self._recv_chunk(chunk_length, timeout)
            except (TimeoutError, ConnectionError) as err:
                logger.error(f"receiving data payload due to {str(err)} error, retrying")
                self.retry_times += 1
                continue

            picked_raw_data += raw_pack
            self._send_chunk(self.ACK_RESPONSE)
            current_pack_idx += 1

            if pack_index == pack_cnt:
                break

            if not self.is_connected:
                self.is_connected = True

        telemetry_crc32 = int.from_bytes(picked_raw_data[:self.CRC32_LENGTH], self.code_order)
        data_payload = picked_raw_data[self.CRC32_LENGTH:]
        if telemetry_crc32 != zlib.crc32(data_payload):
            logger.error(f"crc32 check failed telemetry value {telemetry_crc32} cpr {zlib.crc32(data_payload)}"
                         f", drop frame")
            self.retry_times += 1
            self._send_chunk(b'\x6e')
            return {}

        self.current_link_in_speed = (len(picked_raw_data) +
                                      self.PACK_INDEX_LENGTH * 2 +
                                      self.FRAME_LENGTH) / (time.time() - start_time)
        return json.loads(zlib.decompress(data_payload).decode(self.json_str_coding))

    def close(self):
        """关闭连接"""
        pass

    def _send_chunk(self, _b: bytes) -> int:
        """
        发送数据，通讯实现重载此函数，此函数不能抛出异常，发送失败使用 0 返回值
        :param _b: 请求发送字节长度
        :return: 发送的字节长度，发送失败返回 0
        """
        pass

    def _recv_chunk(self, _n: int, _timeout: float) -> bytes:
        """
        接收数据，使用 recv_timeout 属性作为超时时间
        :param _n: 最大字节长度
        :param _timeout: 传输超时
        :raise TimeoutError: 传输超时
        :raise ConnectionError: 传输错误
        :return: 接收到的原始字节
        """
        pass


class UnixSocketRRT(_BaseRemoteRequestTelemetry):
    """
    Unix 域套接字服务端，同时作为所有网络后端的基类，其他网络后端只重新实现套接字绑定流程，
    网络后端是推荐的 MJ RPC 连接方式，同时也用于大流量的图像传输连接。
    """

    def __init__(self, socket_path: str, **kwargs):
        """
        :param socket_path: 本地 Unix 域套接字地址
        :keyword is_server: 指定是否为服务器，默认 false
        """
        self.is_server: bool = kwargs.pop("is_server", False)
        _BaseRemoteRequestTelemetry.__init__(self, **kwargs)

        self.socket_path = socket_path
        self.tele_socket: socket.socket | None = None

        if self.is_server:
            if os.path.exists(self.socket_path):
                os.remove(self.socket_path)
            try:
                self.server_socket = socket.create_server(self.socket_path,
                                                          backlog=1,
                                                          family=socket.AF_UNIX)
            except OSError as err:
                raise ConnectionError(err)
        else:
            try:
                self.tele_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
                self.tele_socket.connect(self.socket_path)
            except (TimeoutError, ConnectionRefusedError, InterruptedError) as err:
                raise ConnectionError(err)

    def close(self):
        if self.tele_socket:
            self.tele_socket.close()
        if self.is_server and self.server_socket:
            self.server_socket.close()

        super().close()

    def _recv_chunk(self, _n: int, _timeout: float) -> bytes:
        # 服务器模式下如果还没有连接，先等待客户端连接
        if self.is_server and not self.tele_socket:
            self.tele_socket, client_address = self.server_socket.accept()
        self.tele_socket.settimeout(_timeout)

        buf = bytearray()
        while len(buf) < _n:
            try:
                chunk = self.tele_socket.recv(_n - len(buf))
            except (BrokenPipeError, OSError, ConnectionError, socket.timeout) as err:
                raise ConnectionError(err)

            if not chunk:
                raise ConnectionError("socket peer drain, disconnected")
            else:
                buf.extend(chunk)

        return bytes(buf)

    def _send_chunk(self, _b: bytes) -> int:
        try:
            return self.tele_socket.send(_b)
        except (BrokenPipeError, OSError, ConnectionError) as err:
            raise ConnectionError(err)


class IPv4SocketRRT(UnixSocketRRT):
    """
    调整地址解析方式创建 IPV4 连接，主要用于通过网桥或者局域网连接的飞行器。
    """

    def __init__(self, address: str, **kwargs):
        """
        :param address: IPV4 服务器地址 0.0.0.0:6190
        :keyword is_server: 指定是否为服务器，默认 false
        """
        self.is_server: bool = kwargs.pop("is_server", False)
        _BaseRemoteRequestTelemetry.__init__(self, **kwargs)

        addr, str_port = address.split(":")
        self.v4_addr = (addr, int(str_port))
        self.tele_socket: socket.socket | None = None

        if self.is_server:
            try:
                self.server_socket = socket.create_server(self.v4_addr,
                                                          backlog=1,
                                                          family=socket.AF_INET,
                                                          reuse_port=True)
            except OSError as err:
                raise ConnectionError(err)
        else:
            try:
                self.tele_socket = socket.create_connection(self.v4_addr,
                                                            timeout=self.transmit_timeout)
            except (TimeoutError, ConnectionRefusedError, InterruptedError) as err:
                raise ConnectionError(err)


class SerialRRT(_BaseRemoteRequestTelemetry):
    """
    串口遥测后端，可用于 MJ RPC 传输，但不建议用于大流量传输。
    """

    def __init__(self, serial_port: str, **kwargs):
        """
        :param serial_port: 串口设备地址
        :keyword baudrate: 串口传输波特率，默认 57600
        """
        self.baudrate: int = kwargs.pop("baudrate", 57600)
        _BaseRemoteRequestTelemetry.__init__(self, **kwargs)

        self.port = serial_port
        self.tele_serial: serial.Serial | None = None

        try:
            self.tele_serial = serial.Serial(serial_port, self.baudrate, timeout=self.transmit_timeout)
        except serial.SerialException as err:
            raise ConnectionError(err)

    def close(self):
        if self.tele_serial:
            self.tele_serial.close()

        super().close()

    def _send_chunk(self, _b: bytes) -> int:
        try:
            return self.tele_serial.write(_b)
        except serial.serialutil.SerialException as err:
            raise ConnectionError(err)

    def _recv_chunk(self, _n: int, _timeout: float) -> bytes:
        self.tele_serial.timeout = _timeout

        buf = bytearray()

        while len(buf) < _n:
            try:
                chunk = self.tele_serial.read(_n - len(buf))
            except serial.serialutil.SerialException as err:
                raise ConnectionError(err)

            if not chunk:
                raise ConnectionError("")
            buf.extend(chunk)
        return bytes(buf)
